#include "sha512.h"
#include "comm.h"
#include <string.h>
#include <assert.h>

#define SHA512_H0 0x6a09e667f3bcc908ULL
#define SHA512_H1 0xbb67ae8584caa73bULL
#define SHA512_H2 0x3c6ef372fe94f82bULL
#define SHA512_H3 0xa54ff53a5f1d36f1ULL
#define SHA512_H4 0x510e527fade682d1ULL
#define SHA512_H5 0x9b05688c2b3e6c1fULL
#define SHA512_H6 0x1f83d9abfb41bd6bULL
#define SHA512_H7 0x5be0cd19137e2179ULL
#define SHA512_DATA_LENGTH 16

#define MD_INCR(ctx) ((ctx)->count_high += !++(ctx)->count_low)

#define S0(x) (ROTL64(36, (x)) ^ ROTL64(30, (x)) ^ ROTL64(25, (x)))
#define S1(x) (ROTL64(50, (x)) ^ ROTL64(46, (x)) ^ ROTL64(23, (x)))

#define s0(x) (ROTL64(63, (x)) ^ ROTL64(56, (x)) ^ ((x) >> 7))
#define s1(x) (ROTL64(45, (x)) ^ ROTL64(3, (x)) ^ ((x) >> 6))

#define EXPAND(W, i) \
    (W[(i) & 15] +=  \
     (s1(W[((i) - 2) & 15]) + W[((i) - 7) & 15] + s0(W[((i) - 15) & 15])))

#define ROUND(a, b, c, d, e, f, g, h, k, data)   \
    do                                           \
    {                                            \
        h += S1(e) + Choice(e, f, g) + k + data; \
        d += h;                                  \
        h += S0(a) + Majority(a, b, c);          \
    } while (0)

static const uint64_t K[80] = {
    0x428A2F98D728AE22ULL,
    0x7137449123EF65CDULL,
    0xB5C0FBCFEC4D3B2FULL,
    0xE9B5DBA58189DBBCULL,
    0x3956C25BF348B538ULL,
    0x59F111F1B605D019ULL,
    0x923F82A4AF194F9BULL,
    0xAB1C5ED5DA6D8118ULL,
    0xD807AA98A3030242ULL,
    0x12835B0145706FBEULL,
    0x243185BE4EE4B28CULL,
    0x550C7DC3D5FFB4E2ULL,
    0x72BE5D74F27B896FULL,
    0x80DEB1FE3B1696B1ULL,
    0x9BDC06A725C71235ULL,
    0xC19BF174CF692694ULL,
    0xE49B69C19EF14AD2ULL,
    0xEFBE4786384F25E3ULL,
    0x0FC19DC68B8CD5B5ULL,
    0x240CA1CC77AC9C65ULL,
    0x2DE92C6F592B0275ULL,
    0x4A7484AA6EA6E483ULL,
    0x5CB0A9DCBD41FBD4ULL,
    0x76F988DA831153B5ULL,
    0x983E5152EE66DFABULL,
    0xA831C66D2DB43210ULL,
    0xB00327C898FB213FULL,
    0xBF597FC7BEEF0EE4ULL,
    0xC6E00BF33DA88FC2ULL,
    0xD5A79147930AA725ULL,
    0x06CA6351E003826FULL,
    0x142929670A0E6E70ULL,
    0x27B70A8546D22FFCULL,
    0x2E1B21385C26C926ULL,
    0x4D2C6DFC5AC42AEDULL,
    0x53380D139D95B3DFULL,
    0x650A73548BAF63DEULL,
    0x766A0ABB3C77B2A8ULL,
    0x81C2C92E47EDAEE6ULL,
    0x92722C851482353BULL,
    0xA2BFE8A14CF10364ULL,
    0xA81A664BBC423001ULL,
    0xC24B8B70D0F89791ULL,
    0xC76C51A30654BE30ULL,
    0xD192E819D6EF5218ULL,
    0xD69906245565A910ULL,
    0xF40E35855771202AULL,
    0x106AA07032BBD1B8ULL,
    0x19A4C116B8D2D0C8ULL,
    0x1E376C085141AB53ULL,
    0x2748774CDF8EEB99ULL,
    0x34B0BCB5E19B48A8ULL,
    0x391C0CB3C5C95A63ULL,
    0x4ED8AA4AE3418ACBULL,
    0x5B9CCA4F7763E373ULL,
    0x682E6FF3D6B2B8A3ULL,
    0x748F82EE5DEFB2FCULL,
    0x78A5636F43172F60ULL,
    0x84C87814A1F0AB72ULL,
    0x8CC702081A6439ECULL,
    0x90BEFFFA23631E28ULL,
    0xA4506CEBDE82BDE9ULL,
    0xBEF9A3F7B2C67915ULL,
    0xC67178F2E372532BULL,
    0xCA273ECEEA26619CULL,
    0xD186B8C721C0C207ULL,
    0xEADA7DD6CDE0EB1EULL,
    0xF57D4F7FEE6ED178ULL,
    0x06F067AA72176FBAULL,
    0x0A637DC5A2C898A6ULL,
    0x113F9804BEF90DAEULL,
    0x1B710B35131C471BULL,
    0x28DB77F523047D84ULL,
    0x32CAAB7B40C72493ULL,
    0x3C9EBE0A15C9BEBCULL,
    0x431D67C49C100D4CULL,
    0x4CC5D4BECB3E42B6ULL,
    0x597F299CFC657E2AULL,
    0x5FCB6FAB3AD6FAECULL,
    0x6C44198C4A475817ULL,
};

void _sha512_compress(uint64_t *state, const uint8_t *input, const uint64_t *k)
{
    uint64_t data[SHA512_DATA_LENGTH];
    uint64_t A, B, C, D, E, F, G, H; /* Local vars */
    unsigned i;
    uint64_t *d;

    for (i = 0; i < SHA512_DATA_LENGTH; i++, input += 8)
    {
        data[i] = GET_UINT64(input);
    }

    /* Set up first buffer and local data buffer */
    A = state[0];
    B = state[1];
    C = state[2];
    D = state[3];
    E = state[4];
    F = state[5];
    G = state[6];
    H = state[7];

    /* Heavy mangling */
    /* First 16 subrounds that act on the original data */

    for (i = 0, d = data; i < 16; i += 8, k += 8, d += 8)
    {
        ROUND(A, B, C, D, E, F, G, H, k[0], d[0]);
        ROUND(H, A, B, C, D, E, F, G, k[1], d[1]);
        ROUND(G, H, A, B, C, D, E, F, k[2], d[2]);
        ROUND(F, G, H, A, B, C, D, E, k[3], d[3]);
        ROUND(E, F, G, H, A, B, C, D, k[4], d[4]);
        ROUND(D, E, F, G, H, A, B, C, k[5], d[5]);
        ROUND(C, D, E, F, G, H, A, B, k[6], d[6]);
        ROUND(B, C, D, E, F, G, H, A, k[7], d[7]);
    }

    for (; i < 80; i += 16, k += 16)
    {
        ROUND(A, B, C, D, E, F, G, H, k[0], EXPAND(data, 0));
        ROUND(H, A, B, C, D, E, F, G, k[1], EXPAND(data, 1));
        ROUND(G, H, A, B, C, D, E, F, k[2], EXPAND(data, 2));
        ROUND(F, G, H, A, B, C, D, E, k[3], EXPAND(data, 3));
        ROUND(E, F, G, H, A, B, C, D, k[4], EXPAND(data, 4));
        ROUND(D, E, F, G, H, A, B, C, k[5], EXPAND(data, 5));
        ROUND(C, D, E, F, G, H, A, B, k[6], EXPAND(data, 6));
        ROUND(B, C, D, E, F, G, H, A, k[7], EXPAND(data, 7));
        ROUND(A, B, C, D, E, F, G, H, k[8], EXPAND(data, 8));
        ROUND(H, A, B, C, D, E, F, G, k[9], EXPAND(data, 9));
        ROUND(G, H, A, B, C, D, E, F, k[10], EXPAND(data, 10));
        ROUND(F, G, H, A, B, C, D, E, k[11], EXPAND(data, 11));
        ROUND(E, F, G, H, A, B, C, D, k[12], EXPAND(data, 12));
        ROUND(D, E, F, G, H, A, B, C, k[13], EXPAND(data, 13));
        ROUND(C, D, E, F, G, H, A, B, k[14], EXPAND(data, 14));
        ROUND(B, C, D, E, F, G, H, A, k[15], EXPAND(data, 15));
    }

    /* Update state */
    state[0] += A;
    state[1] += B;
    state[2] += C;
    state[3] += D;
    state[4] += E;
    state[5] += F;
    state[6] += G;
    state[7] += H;
}

#define SHA512_COMPRESS(ctx, data) (_sha512_compress((ctx)->state, (data), K))

void sha512_init(struct sha512_ctx *ctx)
{
    ctx->state[0] = SHA512_H0;
    ctx->state[1] = SHA512_H1;
    ctx->state[2] = SHA512_H2;
    ctx->state[3] = SHA512_H3;
    ctx->state[4] = SHA512_H4;
    ctx->state[5] = SHA512_H5;
    ctx->state[6] = SHA512_H6;
    ctx->state[7] = SHA512_H7;
    ctx->count_high = 0;
    ctx->count_low = 0;
    ctx->index = 0;
}

void sha512_update(struct sha512_ctx *ctx, const uint8_t *data, size_t length)
{
    MD_UPDATE(ctx, length, data, SHA512_COMPRESS, MD_INCR(ctx));
}

static void sha512_write_digest(
    struct sha512_ctx *ctx, size_t length, uint8_t *digest)
{
    uint64_t high, low;

    unsigned i;
    unsigned words;
    unsigned leftover;

    assert(length <= SHA512_DIGEST_SIZE);

    MD_PAD(ctx, 16, SHA512_COMPRESS);

    /* There are 1024 = 2^10 bits in one block */
    high = (ctx->count_high << 10) | (ctx->count_low >> 54);
    low = (ctx->count_low << 10) | (ctx->index << 3);

    /* This is slightly inefficient, as the numbers are converted to
       big-endian format, and will be converted back by the compression
       function. It's probably not worth the effort to fix this. */
    PUT_UINT64(ctx->block + (SHA512_BLOCK_SIZE - 16), high);
    PUT_UINT64(ctx->block + (SHA512_BLOCK_SIZE - 8), low);
    SHA512_COMPRESS(ctx, ctx->block);

    words = length / 8;
    leftover = length % 8;

    for (i = 0; i < words; i++, digest += 8)
        PUT_UINT64(digest, ctx->state[i]);

    if (leftover)
    {
        /* Truncate to the right size */
        uint64_t word = ctx->state[i] >> (8 * (8 - leftover));

        do
        {
            digest[--leftover] = word & 0xff;
            word >>= 8;
        } while (leftover);
    }
}

void sha512_final(struct sha512_ctx *ctx, uint8_t *digest, size_t length)
{
    sha512_write_digest(ctx, length, digest);
}
